import os
import csv
import sys
import re
from surprise import Dataset
from surprise import Reader
from collections import defaultdict
import numpy as np


class MovieLens:
    movieID_to_name = {}
    name_to_movieID = {}
    ratingsPath = '../ml-latest-small/ratings.csv'
    moviesPath = '../ml-latest-small/movies.csv'

    def loadMovieLensLatestSmall(self):
        '''
        加载 MovieLens 最新小型数据集，编码为电影名和电影ID 的对应字典
        :return: Surprise 库的 Dataset 对象
        '''
        # 切换到脚本所在的目录
        os.chdir(os.path.dirname(sys.argv[0]))

        ratingsDataset = 0
        self.movieID_to_name = {}
        self.name_to_movieID = {}

        # 定义 Surprise 库的 Reader，指定数据格式
        reader = Reader(line_format='user item rating timestamp', sep=',', skip_lines=1)

        # 使用 Surprise 库加载评分数据
        ratingsDataset = Dataset.load_from_file(self.ratingsPath, reader=reader)

        # 读取电影数据文件，构建电影ID到电影名和电影名到电影ID的映射
        with open(self.moviesPath, newline='', encoding='ISO-8859-1') as csvfile:
            movieReader = csv.reader(csvfile)
            next(movieReader)  # 跳过表头行
            for row in movieReader:
                movieID = int(row[0])
                movieName = row[1]
                self.movieID_to_name[movieID] = movieName
                self.name_to_movieID[movieName] = movieID

        return ratingsDataset

    def getUserRatings(self, user):
        '''
        获取指定用户的评分数据
        :param user: 用户ID
        :return: 用户的评分列表，每个元素为 (movieID, rating)
        '''
        userRatings = []
        hitUser = False
        with open(self.ratingsPath, newline='') as csvfile:
            ratingReader = csv.reader(csvfile)
            next(ratingReader)  # 跳过表头行
            for row in ratingReader:
                userID = int(row[0])
                if (user == userID):
                    movieID = int(row[1])
                    rating = float(row[2])
                    userRatings.append((movieID, rating))
                    hitUser = True
                if (hitUser and (user != userID)):
                    break  # 如果已经读取到其他用户的评分，停止读取

        return userRatings

    def getPopularityRanks(self):
        '''
        获取电影的流行度排名
        :return: 电影ID到流行度排名的字典
        '''
        ratings = defaultdict(int)
        rankings = defaultdict(int)
        with open(self.ratingsPath, newline='') as csvfile:
            ratingReader = csv.reader(csvfile)
            next(ratingReader)  # 跳过表头行
            for row in ratingReader:
                movieID = int(row[1])
                ratings[movieID] += 1  # 统计每部电影的评分次数
        rank = 1
        # 根据评分次数降序排序，生成排名
        # 根据 rating[1](movie出现次数) 排序, 然后遍历, 将出现次数改为 rank(排名)
        for movieID, ratingCount in sorted(ratings.items(), key=lambda x: x[1], reverse=True):
            rankings[movieID] = rank  # 设置排名
            rank += 1  # 下一个的排名变为当前rank+1
        return rankings

    def getGenres(self):
        '''
        获取电影的标签信息
        :return: 电影ID到标签ID列表的字典
        '''
        genres = defaultdict(list)
        genreIDs = {}  # 标签到ID的字典
        maxGenreID = 0  # 当前最大的标签ID号
        with open(self.moviesPath, newline='', encoding='ISO-8859-1') as csvfile:
            movieReader = csv.reader(csvfile)
            next(movieReader)  # 跳过表头行
            for row in movieReader:
                movieID = int(row[0])
                genreList = row[2].split('|')  # 分割标签字符串
                genreIDList = []
                # 将文本形式的标签list, 转为ID数字形式
                for genre in genreList:
                    # 如果已经编码在ID里
                    if genre in genreIDs:
                        # 使用已经编码的ID
                        genreID = genreIDs[genre]
                    else:
                        # 给标签编码
                        genreID = maxGenreID
                        genreIDs[genre] = genreID  # 给新标签分配新的ID
                        maxGenreID += 1  # maxId+1, 避免重复
                    genreIDList.append(genreID)
                genres[movieID] = genreIDList
        # 将标签ID列表转换为位向量
        # 将每个电影的标签ID列表转换为一个二进制向量，其中每个位置表示一个特定的标签是否存在。
        # Convert integer-encoded genre lists to bitfields that we can treat as vectors
        for (movieID, genreIDList) in genres.items():
            # 创建一个长度为 maxGenreID 的列表 bitfield，初始值为0
            bitfield = [0] * maxGenreID
            # 统计id次数
            for genreID in genreIDList:
                # 对于每个标签ID genreID，将 bitfield 中对应位置的值设置为1
                # 表示该标签存在于该电影中。
                bitfield[genreID] = 1
            # 将转换后的位向量 bitfield 存储回 genres 字典中，键为 movieID。
            genres[movieID] = bitfield

        return genres

    def getYears(self):
        '''
        获取电影的上映年份
        :return: 电影ID到上映年份的字典
        '''
        # 匹配四位数字
        p = re.compile(r"(?:\((\d{4})\))?\s*$")
        years = defaultdict(int)
        with open(self.moviesPath, newline='', encoding='ISO-8859-1') as csvfile:
            movieReader = csv.reader(csvfile)
            next(movieReader)  # 跳过表头行
            for row in movieReader:
                movieID = int(row[0])
                title = row[1]  # 电影标题，例如: Toy Story (1995)
                m = p.search(title)  # 使用正则匹配年份
                year = m.group(1)
                if year:
                    years[movieID] = int(year)  # 存储电影ID和年份
        return years

    def getMiseEnScene(self):
        '''
        获取电影的视觉特征数据
        :return: 电影ID到视觉特征列表的字典
        '''
        mes = defaultdict(list)
        with open("LLVisualFeatures13K_Log.csv", newline='') as csvfile:
            mesReader = csv.reader(csvfile)
            next(mesReader)  # 跳过表头行
            for row in mesReader:
                movieID = int(row[0])
                avgShotLength = float(row[1])
                meanColorVariance = float(row[2])
                stddevColorVariance = float(row[3])
                meanMotion = float(row[4])
                stddevMotion = float(row[5])
                meanLightingKey = float(row[6])
                numShots = float(row[7])
                mes[movieID] = [avgShotLength, meanColorVariance, stddevColorVariance,
                                meanMotion, stddevMotion,
                                meanLightingKey, numShots]
        return mes

    def getMovieName(self, movieID):
        '''
        根据电影ID获取电影名
        :param movieID: 电影ID
        :return: 电影名，如果不存在则返回空字符串
        '''
        if movieID in self.movieID_to_name:
            return self.movieID_to_name[movieID]
        else:
            return ""

    def getMovieID(self, movieName):
        '''
        根据电影名获取电影ID
        :param movieName: 电影名
        :return: 电影ID，如果不存在则返回0
        '''
        if movieName in self.name_to_movieID:
            return self.name_to_movieID[movieName]
        else:
            return 0
